import torch

def comp_whisper_ckp(ckp_path1, ckp_path2):
    # Load the checkpoints
    ckp1 = torch.load(ckp_path1, map_location='cpu')
    ckp2 = torch.load(ckp_path2, map_location='cpu')
    
    # Extract model keys if 'model' exists in both checkpoints
    if 'model' in ckp1 and 'model' in ckp2:
        keys1 = set(ckp1['model'].keys())
        keys2 = set(ckp2['model'].keys())
    else:
        print("'model' key not found in one or both checkpoints.")
        return
    
    # Find common and unique keys
    common_keys = keys1.intersection(keys2)
    unique_keys_1 = keys1 - keys2
    unique_keys_2 = keys2 - keys1
    
    # Print results
    # print(f"Keys in both checkpoints ({len(common_keys)}): {common_keys}\n")
    print(f"Keys only in {ckp_path1} ({len(unique_keys_1)}): {unique_keys_1}\n")
    print(f"Keys only in {ckp_path2} ({len(unique_keys_2)}): {unique_keys_2}\n")

# Paths to checkpoints
ckp_path1 = "tmp_whisper/iter_0000001/mp_rank_00/model_optim_rng.pt"
ckp_path2 = "whisper-large-v3-megatron-TP1-TE/iter_0000001/mp_rank_00/model_optim_rng.pt"

# Compare checkpoints
comp_whisper_ckp(ckp_path1, ckp_path2)